import torch
import torch.nn.functional as F
from tqdm import tqdm

from loss import get_loss


@torch.inference_mode()
def evaluate(net, dataloader, device, amp):
    net.eval()
    num_val_batches = len(dataloader)
    val_loss = 0

    # iterate over the validation set
    with torch.autocast(device.type if device.type != 'mps' else 'cpu', enabled=amp):
        for batch in tqdm(dataloader, total=num_val_batches, desc='Validation round', unit='batch', leave=False):
            image, mask_true = batch['image'], batch['mask']

            # move images and labels to correct device and type
            image = image.to(device=device, dtype=torch.float32, memory_format=torch.channels_last)
            mask_true = mask_true.to(device=device, dtype=torch.long)

            # predict the mask
            mask_pred = net(image)

            assert mask_true.min() >= 0 and mask_true.max() <= 1, 'True mask indices should be in [0, 1]'
            # mask_pred = (F.sigmoid(mask_pred) > 0.5).float()
            # compute the loss
            dist_loss, balance_loss, label_loss = get_loss(mask_pred, mask_true)
            val_loss = dist_loss + balance_loss + label_loss

            print("")


    net.train()
    return val_loss / max(num_val_batches, 1)
